// Copyright 2022 Chen Jun
// Licensed under the MIT License.

#include "armor_detector/detector_node.hpp"

#include <cv_bridge/cv_bridge.h>
#include <rmw/qos_profiles.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/convert.h>

#include <ament_index_cpp/get_package_share_directory.hpp>
#include <image_transport/image_transport.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <rclcpp/duration.hpp>
#include <rclcpp/qos.hpp>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

// STD
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <vector>

namespace rm_auto_aim
{
ArmorDetectorNode::ArmorDetectorNode(const rclcpp::NodeOptions & options)
: Node("armor_detector", options)
{
  RCLCPP_INFO(this->get_logger(), "Starting DetectorNode!");

  yaml = YAML::LoadFile("./config/armor_detector.yaml");

  /* ------------------ Detector init ---------------------- */
  detector_ = initDetector();
  RCLCPP_INFO(this->get_logger(), "Detector inited!");
  /* ------------------ Publisher init ------------------------- */
  armors_pub_ = this->create_publisher<detector2tracker::msg::Armors>(
    "/detector/armors", rclcpp::SensorDataQoS());

  /*  ------------------- Camera init ---------------------------- */
  cam_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
    "/camera_info", rclcpp::SensorDataQoS(),
    [this](sensor_msgs::msg::CameraInfo::ConstSharedPtr camera_info) {
      cam_center_ = cv::Point2f(camera_info->k[2], camera_info->k[5]);
      cam_info_ = std::make_shared<sensor_msgs::msg::CameraInfo>(*camera_info);
      pnp_solver_ = std::make_unique<PnPSolver>(camera_info->k, camera_info->d);
      cam_info_sub_.reset();
    });

  img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
    "/image_raw", rclcpp::SensorDataQoS(),
    std::bind(&ArmorDetectorNode::imageCallback, this, std::placeholders::_1));

  /*----------------------debug init --------------------------------*/
  debug_ = declare_parameter("debug", false);
  if(debug_){
    createDebugPublisher();
  }
  debug_param_handler_ = std::make_shared<rclcpp::ParameterEventHandler>(this);
  debug_cb_handle_ =
    debug_param_handler_->add_parameter_callback("debug", [this](const rclcpp::Parameter & param) {
      debug_ = param.as_bool();
      if (debug_) {
        createDebugPublisher();
      } else {
        destroyDebugPublisher();
      }
    });
}

void ArmorDetectorNode::createDebugPublisher()
{
  debug_img_pub_ =
    image_transport::create_publisher(this, "debug_detector_img", rmw_qos_profile_default);
  
  marker_pub_ =
    this->create_publisher<visualization_msgs::msg::MarkerArray>("/detector/marker", 10);

  armor_marker_.ns = "armors";
  armor_marker_.action = visualization_msgs::msg::Marker::ADD;
  armor_marker_.type = visualization_msgs::msg::Marker::CUBE;
  armor_marker_.scale.x = 0.05;
  armor_marker_.scale.z = 0.125;
  armor_marker_.color.a = 1.0;
  armor_marker_.color.g = 0.5;
  armor_marker_.color.b = 1.0;
  armor_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

  text_marker_.ns = "classification";
  text_marker_.action = visualization_msgs::msg::Marker::ADD;
  text_marker_.type = visualization_msgs::msg::Marker::TEXT_VIEW_FACING;
  text_marker_.scale.z = 0.1;
  text_marker_.color.a = 1.0;
  text_marker_.color.r = 1.0;
  text_marker_.color.g = 1.0;
  text_marker_.color.b = 1.0;
  text_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);
  
}

void ArmorDetectorNode::destroyDebugPublisher() 
{ 
  debug_img_pub_.shutdown(); 
  marker_pub_.reset();
}

Eigen::Vector3d rotationMatrixToEulerAngles(Eigen::Matrix3d &R) {
    double sy = sqrt(R(0, 0) * R(0, 0) + R(1, 0) * R(1, 0));
    bool singular = sy < 1e-6;
    double x, y, z;
    if(!singular) {
        x = atan2(R(2, 1), R(2, 2));
        y = atan2(-R(2, 0), sy);
        z = atan2(R(1, 0), R(0, 0));
    } else {
        x = atan2(-R(1, 2), R(1, 1));
        y = atan2(-R(2, 0), sy);
        z = 0;
    }
    return {z, y, x};
}


void ArmorDetectorNode::imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr img_msg)
{
  auto img = cv_bridge::toCvShare(img_msg, "bgr8")->image;
  auto armors = detector_->detect(img);

  if (pnp_solver_ != nullptr) {
    armors_msg_.header = armor_marker_.header = text_marker_.header = img_msg->header;
    armors_msg_.armors.clear();
    marker_array_.markers.clear();
    armor_marker_.id = 0;
    text_marker_.id = 1;

    detector2tracker::msg::Armor armor_msg;
    for (const auto & armor : armors) {
      cv::Mat rvec, tvec;
      bool success = pnp_solver_->solvePnP(armor, rvec, tvec);
      if (success) {
        createMessage(armor_msg,armor,rvec,tvec);
        if(debug_)
          createMarker(armor,armor_msg);
        armors_msg_.armors.emplace_back(armor_msg);

      } else {
        RCLCPP_WARN(this->get_logger(), "PnP failed!");
      }
    }

    if (debug_) {
      for (const auto & armor : armors) {
        detector_->DrawResult(img, armor.id_, armor.confidence_, armor.armor_verticles_);
      }
      debug_img_pub_.publish(cv_bridge::CvImage(img_msg->header, "bgr8", img).toImageMsg());
    }
    // Publishing detected armors
    armors_pub_->publish(armors_msg_);

    // Publishing marker 
    if(debug_)
      publishMarkers();
  }
}

void ArmorDetectorNode::createMessage(detector2tracker::msg::Armor &armor_msg,const Armor & armor, const cv::Mat & rvec, const cv::Mat & tvec)
{
  // Fill basic info
  armor_msg.id = armor.id_;
  armor_msg.type = (int)armor.type;

  // Fill pose
  armor_msg.pose.position.x = tvec.at<double>(0);
  armor_msg.pose.position.y = tvec.at<double>(1);
  armor_msg.pose.position.z = tvec.at<double>(2);

  cv::Mat rotation_matrix;
  cv::Rodrigues(rvec, rotation_matrix);

  tf2::Matrix3x3 tf2_rotation_matrix(
    rotation_matrix.at<double>(0, 0), rotation_matrix.at<double>(0, 1),
    rotation_matrix.at<double>(0, 2), rotation_matrix.at<double>(1, 0),
    rotation_matrix.at<double>(1, 1), rotation_matrix.at<double>(1, 2),
    rotation_matrix.at<double>(2, 0), rotation_matrix.at<double>(2, 1),
    rotation_matrix.at<double>(2, 2));
  tf2::Quaternion tf2_q;
  tf2_rotation_matrix.getRotation(tf2_q);
  armor_msg.pose.orientation = tf2::toMsg(tf2_q);

  // Fill the distance to image center
  armor_msg.distance_to_image_center = pnp_solver_->calculateDistanceToCenter(armor.center);
}

void ArmorDetectorNode::createMarker(const Armor & armor,const detector2tracker::msg::Armor &armor_msg){
  armor_marker_.id++;
  armor_marker_.scale.y = (armor.type == ArmorType::SMALL)
                            ? 0.135
                            : 0.23;

  armor_marker_.pose = armor_msg.pose;
  text_marker_.id++;
  text_marker_.pose.position = armor_msg.pose.position;
  text_marker_.pose.position.y -= 0.1;
  text_marker_.text = class2name[armor.id_];

  marker_array_.markers.emplace_back(armor_marker_);
  marker_array_.markers.emplace_back(text_marker_);
}

std::unique_ptr<ArmorDetectorNetwork> ArmorDetectorNode::initDetector()
{
  auto detector = std::make_unique<ArmorDetectorNetwork>(
    yaml["img_size_"].as<int>(), yaml["model_path_"].as<std::string>(),
    yaml["enemy_color_"].as<int>(), yaml["score_threshold_"].as<float>(),
    yaml["nms_threshold_"].as<float>());

  return detector;
}

void ArmorDetectorNode::publishMarkers()
{
  using Marker = visualization_msgs::msg::Marker;
  armor_marker_.action = armors_msg_.armors.empty() ? Marker::DELETE : Marker::ADD;
  marker_array_.markers.emplace_back(armor_marker_);
  marker_pub_->publish(marker_array_);
}

}  // namespace rm_auto_aim

#include "rclcpp_components/register_node_macro.hpp"

// Register the component with class_loader.
// This acts as a sort of entry point, allowing the component to be discoverable when its library
// is being loaded into a running process.
RCLCPP_COMPONENTS_REGISTER_NODE(rm_auto_aim::ArmorDetectorNode)
